import duckdb
import time
import functools
import psutil

duckdb.execute("SET custom_extension_repository = './'")
duckdb.execute("INSTALL spatial;")
duckdb.execute("LOAD spatial;")
# duckdb.install_extension('spatial')
# duckdb.load_extension('spatial')


def calculate_statistics(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        start_memory = psutil.Process().memory_info().rss / 1024 / 1024
        result = func(*args, **kwargs)
        end_time = time.time()
        end_memory = psutil.Process().memory_info().rss / 1024 / 1024
        execution_time = end_time - start_time
        memory_used = end_memory - start_memory
        print(f"函数 {func.__name__} 的运行时间为：{execution_time:.6f} 秒")
        print(f"函数 {func.__name__} 的内存占用为：{memory_used:.6f} MB")
        return result
    return wrapper


@calculate_statistics
def test_duckdb():
    df = duckdb.sql("SELECT * FROM st_read('data/test.xlsx', layer = 'Sheet1', open_options=['HEADERS=FORCE']);").df()
    print(df.head(10))
    print(df.shape)
    print(df.dtypes)


if __name__ == '__main__':
    test_duckdb()
